package websocket

import (
	"bytes"
	"sync"
)

const bucketSize = 4 * 1024

var bufferpool = NewBufferPool()

type BufferPool struct {
	pool sync.Pool
}

func NewBufferPool() *BufferPool {
	var c = new(BufferPool)
	c.pool.New = func() interface{} {
		return bytes.NewBuffer(make([]byte, 0, bucketSize))
	}
	return c
}

func (c *BufferPool) Put(b *bytes.Buffer) {
	if b.Cap() == bucketSize {
		c.pool.Put(b)
	}
}

func (c *BufferPool) Get() *bytes.Buffer {
	buffer := c.pool.Get().(*bytes.Buffer)
	buffer.Reset()
	return buffer
}

type LinkedBuffer struct {
	buffers []*bytes.Buffer
}

func (c *LinkedBuffer) Len() int {
	var sum = 0
	for _, item := range c.buffers {
		sum += item.Len()
	}
	return sum
}

func (c *LinkedBuffer) Gte(n int) bool {
	if n == 0 {
		return true
	}
	var sum = 0
	for _, item := range c.buffers {
		sum += item.Len()
		if sum >= n {
			return true
		}
	}
	return false
}

func (c *LinkedBuffer) getBuffer() *bytes.Buffer {
	var n = len(c.buffers)
	if n == 0 {
		b := bufferpool.Get()
		c.buffers = append(c.buffers, b)
		return b
	}
	return c.buffers[n-1]
}

func (c *LinkedBuffer) Write(p []byte) (int, error) {
	var sum = 0
	var buf = c.getBuffer()

	for {
		var avai = bucketSize - buf.Len()
		if avai >= len(p) {
			num, _ := buf.Write(p)
			sum += num
			return sum, nil
		}

		num, _ := buf.Write(p[:avai])
		sum += num
		p = p[avai:]
		buf = bufferpool.Get()
		c.buffers = append(c.buffers, buf)
	}
}

func (c *LinkedBuffer) pop() {
	for len(c.buffers) > 0 {
		if buf := c.buffers[0]; buf.Len() == 0 {
			c.buffers = c.buffers[1:]
			bufferpool.Put(buf)
		} else {
			return
		}
	}
}

func (c *LinkedBuffer) ZeroCopy(n int) *bytes.Buffer {
	if n == c.Len() && len(c.buffers) == 1 {
		buffer := c.buffers[0]
		c.buffers = c.buffers[1:]
		return buffer
	}
	return nil
}

func (c *LinkedBuffer) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}

	var sum = 0
	for _, buf := range c.buffers {
		if len(p) == 0 {
			c.pop()
			return sum, nil
		}
		num, _ := buf.Read(p)
		sum += num
		p = p[num:]
	}
	c.pop()
	return sum, nil
}

// read http header
func (c *LinkedBuffer) readHeader() ([]byte, bool) {
	if len(c.buffers) == 0 {
		return nil, false
	}

	var buf = c.buffers[0]
	var b = buf.Bytes()
	var n = buf.Len()
	for i := n - 1; i >= 1; i-- {
		if b[i] == '\n' && b[i-1] == '\r' {
			buf.Next(i + 1)
			return b[:i+1], true
		}
	}
	return nil, false
}
